"""
深度CNN模型
使用VGG-like结构的深度卷积网络
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from .common import ConvBlock


class DeepCNNEncoder(nn.Module):
    """
    深度CNN编码器
    VGG-like结构
    """
    def __init__(self, in_channels=1):
        super().__init__()
        
        # VGG-like结构
        self.features = nn.Sequential(
            # Block 1: 1 -> 64
            nn.Conv2d(in_channels, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),  # /2
            
            # Block 2: 64 -> 128
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),  # /4
            
            # Block 3: 128 -> 256
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),  # /8
            
            # Block 4: 256 -> 512
            nn.Conv2d(256, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),  # /16
        )
        
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.output_dim = 512
    
    def forward(self, x):
        """
        前向传播
        
        参数:
            x: [B, 1, H, W]
        
        返回:
            features: [B, 512]
        """
        x = self.features(x)
        x = self.global_pool(x)
        return x.view(x.size(0), -1)


class DeepCNNPredictor(nn.Module):
    """
    深度CNN体积预测模型
    
    使用更深的VGG-like网络提取特征
    相比baseline有更强的表达能力
    
    参数:
        num_images: 每个样本的图像数量（默认20）
        num_features: 传统特征数量（默认6）
    """
    def __init__(self, num_images=20, num_features=6):
        super().__init__()
        self.num_images = num_images
        self.num_features = num_features
        
        # 深度CNN编码器
        self.encoder = DeepCNNEncoder(in_channels=1)
        encoder_dim = self.encoder.output_dim
        
        # 图像特征处理
        self.image_processor = nn.Sequential(
            nn.Linear(encoder_dim, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(inplace=True)
        )
        
        # 像素尺度处理
        self.scale_processor = nn.Sequential(
            nn.Linear(num_images, 32),
            nn.ReLU(inplace=True),
            nn.Linear(32, 16),
            nn.ReLU(inplace=True)
        )
        
        # 传统特征处理
        self.feature_processor = nn.Sequential(
            nn.Linear(num_features, 16),
            nn.ReLU(inplace=True),
            nn.Linear(16, 8),
            nn.ReLU(inplace=True)
        )
        
        # 最终预测层 (128 + 16 + 8 = 152)
        self.predictor = nn.Sequential(
            nn.Linear(152, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.ReLU(inplace=True),
            nn.Linear(64, 1)
        )
    
    def forward(self, images, pixel_scales, traditional_features):
        """
        前向传播
        
        参数:
            images: [B, 20, 1, H, W]
            pixel_scales: [B, 20]
            traditional_features: [B, 6]
        
        返回:
            output: [B, 1]
        """
        batch_size = images.size(0)
        
        # 处理每张图像
        image_features = []
        for i in range(self.num_images):
            img = images[:, i, :, :, :]
            encoded = self.encoder(img)
            image_features.append(encoded)
        
        # 聚合图像特征
        image_features = torch.stack(image_features, dim=1)
        aggregated_features = torch.mean(image_features, dim=1)
        
        # 处理各部分特征
        img_feat = self.image_processor(aggregated_features)
        scale_feat = self.scale_processor(pixel_scales)
        trad_feat = self.feature_processor(traditional_features)
        
        # 组合所有特征
        combined = torch.cat([img_feat, scale_feat, trad_feat], dim=1)
        
        # 预测
        output = self.predictor(combined)
        
        return output
    
    def switch_to_deploy(self):
        """切换到推理模式"""
        self.eval()
        print("Deep CNN model switched to inference mode")


if __name__ == "__main__":
    # 测试模型
    print("Testing Deep CNN Model...")
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = DeepCNNPredictor(num_images=20, num_features=6).to(device)
    
    # 统计参数量
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"\nModel Statistics:")
    print(f"  Total parameters: {total_params:,}")
    print(f"  Trainable parameters: {trainable_params:,}")
    print(f"  Model size: {total_params * 4 / 1024 / 1024:.2f} MB (float32)")
    
    # 测试前向传播
    batch_size = 2
    test_images = torch.randn(batch_size, 20, 1, 96, 288).to(device)
    test_scales = torch.randn(batch_size, 20).to(device)
    test_features = torch.randn(batch_size, 6).to(device)
    
    print(f"\nInput shapes:")
    print(f"  Images: {test_images.shape}")
    print(f"  Scales: {test_scales.shape}")
    print(f"  Features: {test_features.shape}")
    
    with torch.no_grad():
        output = model(test_images, test_scales, test_features)
    
    print(f"\nOutput shape: {output.shape}")
    print(f"Output values: {output.squeeze().cpu().numpy()}")
    
    print("\n✓ Deep CNN model test passed!")